import numpy as np
import pywt


def PQRST(ecg, locatedR, sampling_rate):
    # 输入ECG数据
    ecgdata = ecg

    # 小波变换层数
    level = 7
    # 小波名称
    wavename = 'db8'

    # 进行小波变换
    coeffs = pywt.wavedec(ecgdata, wavename, level=level)
    threshold = np.sqrt(2 * np.log(len(ecgdata)))
    # 去除基线漂移
    coeffs[7] = pywt.threshold(coeffs[7], threshold)  # 假设基线漂移在最后一层
    # 去除高频噪声 1，2层
    for i in range(2):
        coeffs[i] = pywt.threshold(coeffs[i], threshold)

    # 重构信号
    denoised_signal = pywt.waverec(coeffs, wavename)

    '''#######################    各尺度下小波近似系数   各尺度下小波细节系数  ###################
    ##########################################################'''

    level = 4  # 设置小波分解的级别
    signal = denoised_signal
    points = len(signal)

    swa = np.zeros((level, points))  # 存储概貌信息
    swd = np.zeros((level, points))  # 存储细节信息

    # 第一个尺度的信号分解
    for i in range(3, points):
        swa[0, i] = (1 / 4) * signal[i - 2 ** 0 * 0] + (3 / 4) * signal[i - 2 ** 0 * 1] + (3 / 4) * signal[
            i - 2 ** 0 * 2] + (1 / 4) * signal[i - 2 ** 0 * 3]
        swd[0, i] = (-1 / 4) * signal[i - 2 ** 0 * 0] - (3 / 4) * signal[i - 2 ** 0 * 1] + (3 / 4) * signal[
            i - 2 ** 0 * 2] + (1 / 4) * signal[i - 2 ** 0 * 3]

    # 后续尺度的信号分解
    for j in range(1, level):
        for i in range(24, points):
            swa[j, i] = (1 / 4) * swa[j - 1, i - 2 ** j * 0] + (3 / 4) * swa[j - 1, i - 2 ** j * 1] + (3 / 4) * swa[
                j - 1, i - 2 ** j * 2] + (1 / 4) * swa[j - 1, i - 2 ** j * 3]
            swd[j, i] = (-1 / 4) * swa[j - 1, i - 2 ** j * 0] - (3 / 4) * swa[j - 1, i - 2 ** j * 1] + (3 / 4) * swa[
                j - 1, i - 2 ** j * 2] + (1 / 4) * swa[j - 1, i - 2 ** j * 3]

    '''#######################    QRS波检测   ###################
    ##########################################################'''

    # %    求正负极大值对   ###################
    # 初始化数组
    pddw = np.zeros(swd.shape)
    nddw = np.zeros(swd.shape)

    # 小波系数的大于0的点
    posw = swd * (swd > 0)

    # 斜率大于0
    pdw = (posw[:, :-1] - posw[:, 1:]) < 0
    pdw = pdw.astype(int)

    # 正极大值点
    for i in range(1, points - 1):
        pddw[:, i] = (pdw[:, i - 1] - pdw[:, i] > 0).astype(int)

    # 小波系数小于0的点
    negw = swd * (swd < 0)
    ndw = (negw[:, :-1] - negw[:, 1:]) > 0
    ndw = ndw.astype(int)

    # 负极大值点
    for i in range(1, points - 1):
        nddw[:, i] = (ndw[:, i - 1] - ndw[:, i]) > 0

    # 逻辑或运算
    ddw = np.logical_or(pddw, nddw)
    ddw[:, 0] = 1
    ddw[:, -1] = 1

    # 计算峰值点的值，将其他点设置为0
    wpeak = ddw * swd
    wpeak[:, 0] = wpeak[:, 0] + 1e-10
    wpeak[:, -1] = wpeak[:, -1] + 1e-10

    interva2 = np.zeros(points)
    intervaqs = np.zeros(points)
    Mj1 = wpeak[0, :]
    Mj3 = wpeak[2, :]
    Mj4 = wpeak[3, :]

    posi = Mj3 * (Mj3 > 0)

    thposi = (np.max(posi[:points // 4]) + np.max(posi[points // 4:2 * points // 4]) +
              np.max(posi[2 * points // 4:3 * points // 4]) + np.max(posi[3 * points // 4:4 * points // 4])) / 4

    posi = posi > thposi / 3

    nega = Mj3 * (Mj3 < 0)

    thnega = (np.min(nega[:points // 4]) + np.min(nega[points // 4:2 * points // 4]) +
              np.min(nega[2 * points // 4:3 * points // 4]) + np.min(nega[3 * points // 4:4 * points // 4])) / 4

    nega = -1 * (nega < thnega / 4)

    interva = posi + nega

    loca = np.nonzero(interva)[0]

    diff = np.zeros(len(loca))

    for i in range(len(loca) - 1):
        if abs(loca[i] - loca[i + 1]) < 80:
            diff[i] = interva[loca[i]] - interva[loca[i + 1]]
        else:
            diff[i] = 0

    loca2 = np.where(diff == -2)[0]

    interva2 = np.copy(interva)
    interva2[loca[loca2]] = interva[loca[loca2]]
    interva2[loca[loca2 + 1]] = interva[loca[loca2 + 1]]

    intervaqs = interva2[10:points]

    countQ = 0
    countS = 0
    mark1 = 0
    mark2 = 0
    mark3 = 0
    Rnum = 0

    # %*************************求正负极值对过零，即R波峰*******************%

    countR = locatedR
    # if len(denoised_signal) - countR[-1] <= 50:
    #     countR = countR[:-1]
    # if countR[1] <= 50:
    #     countR = countR[1:]

    # 将countQ和countS定义为可变列表
    countQ = [0] * points
    countS = [0] * points

    # 求QRS波起点
    R = countR
    for i in range(len(R)):
        kqs = R[i]
        markq = 0
        while kqs > 0 and markq < 3:
            if Mj1[kqs] != 0:
                markq += 1
            kqs -= 1
        countQ[kqs] = -1

    # 求QRS波终点
    for i in range(len(R)):
        kqs = R[i]
        marks = 0
        while kqs < points - 1 and marks < 3:
            if Mj1[kqs] != 0:
                marks += 1
            kqs += 1
        # print(len(countS))
        # print(R[i])
        # print(kqs)
        countS[kqs] = -1

    '''#######################  基于尺度3的P波检测     ###################
    ##########################################################'''

    countPMj4 = np.zeros(points)
    countTMj4 = np.zeros(points)
    countP = np.zeros(points)
    countPbegin = np.zeros(points)
    countPend = np.zeros(points)
    countT = np.zeros(points)
    countTbegin = np.zeros(points)
    countTend = np.zeros(points)

    windowSize = 100

    Q = [i for i, x in enumerate(countQ) if x != 0]

    for i in range(1, len(R)-1):
        flag = 0
        mark4 = 0
        RRinteral = R[i] - R[i - 1]

        for j in range(0, int((2 / 3) * RRinteral), 5):
            windowEnd = Q[i] - j
            windowBegin = windowEnd - windowSize

            if windowBegin < R[i - 1] + (1 / 3) * RRinteral:
                break

            windowMax = np.max(Mj3[windowBegin:windowEnd])
            windowMin = np.min(Mj3[windowBegin:windowEnd])
            maxindex = np.argmax(Mj3[windowBegin:windowEnd])
            minindex = np.argmin(Mj3[windowBegin:windowEnd])

            if minindex < maxindex and (maxindex - minindex) < (
                    2 / 3) * windowSize and windowMax > 0.01 and windowMin < -0.1:
                flag = 1
                mark4 = round((maxindex + minindex) / 2 + windowBegin)
                countPMj4[mark4] = 1
                countP[mark4 - 17] = -1
                countPbegin[windowBegin + minindex - 17] = -1
                countPend[windowBegin + maxindex - 17] = -1

            if flag == 1:
                break

        if mark4 == 0 and flag == 0:
            mark4 = round(R[i] - (1 / 3) * RRinteral)
            countP[mark4 - 17] = -2  # 未检测出置为-2
            countPbegin[mark4 - 17] = -100
            countPend[mark4 - 17] = -100

    '''#######################  基于尺度3的T波检测     ###################
    ##########################################################'''

    windowSizeQ = 100
    S = [i for i, x in enumerate(countS) if x != 0]

    for i in range(1, len(R) - 1):
        flag = 0
        mark5 = 0
        RRinterval = R[i + 1] - R[i]

        for j in range(0, int(RRinterval * 2 / 3), 5):
            windowBegin = S[i] + int(RRinterval / 8) + j
            windowEnd = windowBegin + windowSizeQ

            if windowEnd > R[i + 1] - RRinterval / 4:
                break

            # Calculate the window's maximum and minimum values
            windowData = Mj3[windowBegin:windowEnd]
            maxindex = np.argmax(windowData)
            minindex = np.argmin(windowData)
            windowMax = np.max(windowData)
            windowMin = np.min(windowData)

            if minindex < maxindex and (maxindex - minindex) < windowSizeQ and windowMax > 0.01 and windowMin < -0.1:
                flag = 1
                mark5 = round((maxindex + minindex) / 2 + windowBegin)
                countTMj4[mark5] = 1
                countT[mark5 - 9] = -1

                # Determine T-wave start and end points
                countTbegin[windowBegin + minindex - 9] = -1
                countTend[windowBegin + maxindex - 9] = -1

            if flag == 1:
                break

        if mark5 == 0:
            mark5 = round(R[i] + RRinterval / 3)
            countT[mark5 - 9] = -2  # 未检测出置为-2
            countTbegin[mark5 - 9] = -100
            countTend[mark5 - 9] = -100
    R = countR
    T = [i for i, x in enumerate(countT) if x != 0]
    P = [i for i, x in enumerate(countP) if x != 0]
    Q = [i for i, x in enumerate(countQ) if x != 0]
    S = [i for i, x in enumerate(countS) if x != 0]
    Pbegin = [i for i, x in enumerate(countPbegin) if x != 0]
    Pend = [i for i, x in enumerate(countPend) if x != 0]
    Tbegin = [i for i, x in enumerate(countTbegin) if x != 0]
    Tend = [i for i, x in enumerate(countTend) if x != 0]
    return R, Q, S, P, T, Pbegin, Pend, Tbegin, Tend
    # return countR, countQ, countS, countP, countT